import torch
import itertools
from utils.image_pool import ImagePool
from .base_model import BaseModel
from .LFNet import networks
from utils.util import SSIM,PSNR,normal,AverageMeter
import numpy as np
from  numpy.fft import fft2,fftshift
# from .lightcnn import network as lightcnn
from models import create_model
# import mypymath
# import time
def get_frequency(src):
            res = []
            for i in range(src.shape[0]):
                res.append(fftshift(fft2(src[i,:,:])))
            return torch.tensor(np.array(res))

class LFNetModel(BaseModel):
    """
    This class implements the CycleGAN model, for learning image-to-image translation without paired data.

    The model training requires '--dataset_mode unaligned' dataset.
    By default, it uses a '--netG resnet_9blocks' ResNet generator,
    a '--netD basic' discriminator (PatchGAN introduced by pix2pix),
    and a least-square GANs objective ('--gan_mode lsgan').

    CycleGAN paper: https://arxiv.org/pdf/1703.10593.pdf
    """
    @staticmethod
    def modify_commandline_options(parser, is_train=True):
        """Add new dataset-specific options, and rewrite default values for existing options.

        Parameters:
            parser          -- original option parser
            is_train (bool) -- whether training phase or test phase. You can use this flag to add training-specific or test-specific options.
        
        Returns:
            the modified parser.

        For CycleGAN, in addition to GAN losses, we introduce lambda_A, lambda_B, and lambda_identity for the following losses.
        A (source domain), B (target domain).
        Generators: G_A: A -> B; G_B: B -> A.
        Discriminators: D_A: G_A(A) vs. B; D_B: G_B(B) vs. A.
        Forward cycle loss:  lambda_A * ||G_B(G_A(A)) - A|| (Eqn. (2) in the paper)
        Backward cycle loss: lambda_B * ||G_A(G_B(B)) - B|| (Eqn. (2) in the paper)
        Identity loss (optional): lambda_identity * (||G_A(B) - B|| * lambda_B + ||G_B(A) - A|| * lambda_A) (Sec 5.2 "Photo generation from paintings" in the paper)
        Dropout is not used in the original CycleGAN paper.
        """
        parser.set_defaults(no_dropout=True)  # default CycleGAN did not use dropout
        parser.set_defaults(frequency=True)
        # parser.add_argument('--num_classes', type=int, default=3, help='the architecture situation of model')
        # parser.add_argument('--pretrained_model', type=str, default="lightcnn", help='the architecture of pretrained model')
        # parser.add_argument('--pretrained_path', type=str, default="lightcnn", help='the path of pretrained model')
        # parser.add_argument('--pretrained_ndf', type=int, default=128, help='the number of features of output ')
        # parser.add_argument('--pretrained_output_nc', type=int, default=256, help='the number of features of output ')
        # parser.add_argument('--bins', type=int, default=16, help='the number of features of output ')
        parser.add_argument('--use_bias', type=bool, default=True, help='use bias in conv blocks or not')
        # parser.add_argument('--trainNet_type', type=str, default='row', help='train row rcnn or col rcnn')
        parser.add_argument('--times_steps', type=int, default=9, help='len of time steps')
        
        # parser.add_argument('--norm', type=str, default=None, help='type of norm layer')
        # parser.add_argument('--frequency', type=bool, default=True, help='the mode of caculating')
        # if is_train:
        # parser.add_argument('--lambda_A', type=float, default=10.0, help='weight for cycle loss (A -> B -> A)')
        # parser.add_argument('--lambda_B', type=float, default=10.0, help='weight for cycle loss (B -> A -> B)')
        # parser.add_argument('--lambda_D', type=float, default=10.0, help='weight for Adversarial network loss')
        # parser.add_argument('--lambda_fre', type=float, default=0.05, help='weight for Frequency loss')
        # parser.add_argument('--lambda_idt', type=float, default=0.5, help='use identity mapping. Setting lambda_identity other than 0 has an effect of scaling the weight of the identity mapping loss. For example, if the weight of the identity loss should be 10 times smaller than the weight of the reconstruction loss, please set lambda_identity = 0.1')
            

        return parser

    def __init__(self, opt):
        """Initialize the CycleGAN class.

        Parameters:
            opt (Option class)-- stores all the experiment flags; needs to be a subclass of BaseOptions
        """
        
        BaseModel.__init__(self, opt)
        # specify the training losses you want to print out. The training/test scripts will call <BaseModel.get_current_losses>
        # self.loss_names = ['D_A', 'G_A', 'cycle_A', 'idt_A', 'D_B', 'G_B', 'cycle_B', 'idt_B']
        # self.loss_chage_line = "\n"

        # define networks (both Generators and discriminators)
        # The naming is different from those used in the paper.
        # Code (vs. paper): G_A (G), G_B (F), D_A (D_Y), D_B (D_X)
        self.visuale_param()
        # if opt.trainNet_type == 'row':
        #     self.model_names = ["lfnetrow",]
        #     self.netlfnetrow = networks.define_G(opt.input_nc, opt.crop_size,opt.times_steps, opt.output_nc, opt.norm,opt.use_bias,opt.init_type, opt.init_gain, self.gpu_ids)
        # elif opt.trainNet_type == 'col':
        #     self.model_names = ["lfnetcol",]
        #     self.netlfnetcol = networks.define_G(opt.input_nc, opt.crop_size,opt.times_steps, opt.output_nc, opt.norm,opt.use_bias,opt.init_type, opt.init_gain, self.gpu_ids)
        # elif opt.trainNet_type == 'ldiag':
        #     self.model_names = ["lfnetldiag",]
        #     self.netlfnetldiag = networks.define_G(opt.input_nc, opt.crop_size,opt.times_steps, opt.output_nc, opt.norm,opt.use_bias,opt.init_type, opt.init_gain, self.gpu_ids)
        # elif opt.trainNet_type == 'rdiag':
        #     self.model_names = ["lfnetrdiag",]
        #     self.netlfnetrdiag = networks.define_G(opt.input_nc, opt.crop_size,opt.times_steps, opt.output_nc, opt.norm,opt.use_bias,opt.init_type, opt.init_gain, self.gpu_ids)
        
        self.model_names = ["lfnet",]
        self.netlfnet = networks.define_G(opt.input_nc, opt.crop_size,opt.times_steps, opt.output_nc, opt.norm,opt.use_bias,opt.init_type, opt.init_gain, self.gpu_ids)
        if self.isTrain:
            # if opt.trainNet_type == 'row':
            #     self.optimizer = torch.optim.Adam( self.netlfnetrow.parameters(),lr=opt.lr, betas=(opt.beta1, 0.999))
            # elif opt.trainNet_type == 'col':
            #     self.optimizer = torch.optim.Adam( self.netlfnetcol.parameters(),lr=opt.lr, betas=(opt.beta1, 0.999))
            # elif opt.trainNet_type == 'ldiag':
            #     self.optimizer = torch.optim.Adam( self.netlfnetldiag.parameters(),lr=opt.lr, betas=(opt.beta1, 0.999))
            # elif opt.trainNet_type == 'rdiag':
            #     self.optimizer = torch.optim.Adam( self.netlfnetrdiag.parameters(),lr=opt.lr, betas=(opt.beta1, 0.999))
            self.optimizer = torch.optim.Adam( self.netlfnet.parameters(),lr=opt.lr, betas=(opt.beta1, 0.999))
                
            self.optimizers.append(self.optimizer)
            self.criterion = torch.nn.MSELoss()
        
    def visuale_param(self):
        self.visual_names = ["pred_hr4","hr_data4","lr_data4"]
        # self.visual_names += ["pred_hr_ldiag4","pred_hr_rdiag4",]
        # self.visual_names += ["hr_data4","lr_data4"]
        # self.model_names += ["lfnetldiag","lfnetrdiag",]
        if self.isTrain:
            # self.loss_names = ["LightField_row","LightField_col","LightField"]
            # self.loss_names = ["mean_LightField_row","mean_LightField_col",
            #                    "mean_LightField_ldiag","mean_LightField_rdiag",
            #                    "mean_LightField"]
            self.loss_names = ["mean_LightField"]#+self.opt.trainNet_type]
            self.count_mean_LightField = AverageMeter()
            # if self.opt.trainNet_type == 'row':
            #     self.count_mean_LightField_row = AverageMeter()
            # elif self.opt.trainNet_type == 'col':
            #     self.count_mean_LightField_col = AverageMeter()
            # elif self.opt.trainNet_type == 'ldiag':
            #     self.count_mean_LightField_ldiag = AverageMeter()
            # elif self.opt.trainNet_type == 'rdiag':
            #     self.count_mean_LightField_rdiag = AverageMeter()
            
        elif not self.isTrain:
            
            self.loss_names = ["mean_hr_SSIM","mean_hr_PSNR"]
            self.count_mean_hr_PSNR = AverageMeter()
            self.count_mean_hr_SSIM = AverageMeter()
            # if self.opt.trainNet_type == 'row':
            #     self.loss_names = ["mean_hr_row_SSIM","mean_hr_row_PSNR"]
            #     self.count_mean_hr_row_PSNR = AverageMeter()
            #     self.count_mean_hr_row_SSIM = AverageMeter()
            # elif self.opt.trainNet_type == 'col':
            #     self.loss_names = ["mean_hr_col_SSIM","mean_hr_col_PSNR",]
            #     self.count_mean_hr_col_PSNR = AverageMeter()
            #     self.count_mean_hr_col_SSIM = AverageMeter()
            # elif self.opt.trainNet_type == 'ldiag':
            #     self.loss_names = ["mean_hr_ldiag_SSIM","mean_hr_ldiag_PSNR",]
            #     self.count_mean_hr_ldiag_PSNR = AverageMeter()
            #     self.count_mean_hr_ldiag_SSIM = AverageMeter()
            # elif self.opt.trainNet_type == 'rdiag':
            #     self.loss_names = ["mean_hr_rdiag_SSIM","mean_hr_rdiag_PSNR",]
            #     self.count_mean_hr_rdiag_PSNR = AverageMeter()
            #     self.count_mean_hr_rdiag_SSIM = AverageMeter()
            
        pass
    
    def set_input(self, input):
        """Unpack input data from the dataloader and perform necessary pre-processing steps.

        Parameters:
            input (dict): include the data itself and its metadata information.

        The option 'direction' can be used to swap domain A and domain B.
        """
        if self.opt.phase == 'row':
            self.lr_data = input['lr_row'].float().to(self.device)
            self.hr_data = input['hr_row'].float().to(self.device)
        elif self.opt.phase == 'col':
            self.lr_data = input['lr_col'].float().to(self.device)
            self.hr_data = input['hr_col'].float().to(self.device)
        elif self.opt.phase == 'ldiag':
            self.lr_data = input['lr_ldiag'].float().to(self.device)
            self.hr_data = input['hr_ldiag'].float().to(self.device)
        elif self.opt.phase == 'rdiag':
            self.lr_data = input['lr_rdiag'].float().to(self.device)
            self.hr_data = input['hr_rdiag'].float().to(self.device)
        
        self.lr_data4 = input['lr_row'][0,4:5,:,:,:]
        self.hr_data4 = input['hr_row'][0,4:5,:,:,:]
        self.image_paths = input['image_path']
        pass
        

    def forward(self):
        """Run forward pass; called by both functions <optimize_parameters> and <test>."""
        
        # if self.opt.trainNet_type == 'row':
        #     self.pred_hr =  self.netlfnetrow(self.lr_data)
        # elif self.opt.trainNet_type == 'col':
        #     self.pred_hr =  self.netlfnetcol(self.lr_data)
        # elif self.opt.trainNet_type == 'ldiag':
        #     self.pred_hr =  self.netlfnetldiag(self.lr_data)
        # elif self.opt.trainNet_type == 'rdiag':
        #     self.pred_hr =  self.netlfnetrdiag(self.lr_data)
        self.pred_hr =  self.netlfnet(self.lr_data)
        self.pred_hr4 = self.pred_hr[0,4:5,:,:,:]
        # self.pred_hr_ldiag4 = self.pred_hr_ldiag[0,4:5,:,:,:]
        # self.pred_hr_rdiag4 = self.pred_hr_rdiag[0,4:5,:,:,:]
        pass
    
    def calculate_parameters(self):

        self.loss_LightField = self.criterion(self.pred_hr,self.hr_data)
        self.loss_LightField.backward()
        
        self.count_mean_LightField.update(self.loss_LightField.contiguous().detach().cpu().item())
        # self.count_mean_LightField_ldiag.update(self.loss_LightField_ldiag.contiguous().detach().cpu().item())
        # self.count_mean_LightField_rdiag.update(self.loss_LightField_rdiag.contiguous().detach().cpu().item())
        # self.count_mean_LightField.update(self.loss_LightField.contiguous().detach().cpu().item())
        
        # self.loss_mean_LightField_row = self.count_mean_LightField_row.avg
        # self.loss_mean_LightField_col = self.count_mean_LightField_col.avg
        # self.loss_mean_LightField_ldiag = self.count_mean_LightField_ldiag.avg
        # self.loss_mean_LightField_rdiag = self.count_mean_LightField_rdiag .avg
        self.loss_mean_LightField = self.count_mean_LightField.avg
        pass
    
    def optimize_parameters(self):
        """Calculate losses, gradients, and update network weights; called in every training iteration"""
        # forward
        self.forward()      # compute fake images and reconstruction images.
        self.optimizer.zero_grad()
        self.calculate_parameters()
        self.optimizer.step()
        
        pass
        
    def cal_score(self):
        self.loss_hr_PSNR = PSNR(self.pred_hr4,self.hr_data4)
        # self.loss_hr_col_PSNR = PSNR(self.pred_hr_col4,self.hr_data4)
        # self.loss_hr_ldiag_PSNR = PSNR(self.pred_ldiag_row4,self.hr_data4)
        # self.loss_hr_rdiag_PSNR = PSNR(self.pred_rdiag_col4,self.hr_data4)

        self.count_mean_hr_PSNR.update(self.loss_hr_PSNR)
        # self.count_mean_hr_col_PSNR.update(self.loss_hr_col_PSNR)
        # self.count_mean_hr_ldiagPSNR.update(self.loss_hr_ldiag_PSNR)
        # self.count_mean_hr_rdiag_PSNR.update(self.loss_hr_rdiag_PSNR)

        self.loss_mean_hr_PSNR = self.count_mean_hr_PSNR.avg
        # self.loss_mean_hr_col_PSNR = self.count_mean_hr_col_PSNR.avg
        # self.loss_mean_hr_ldiag_PSNR = self.count_mean_hr_ldiag_PSNR.avg
        # self.loss_mean_hr_rdiag_PSNR = self.count_mean_hr_rdiag_PSNR.avg
        
        self.loss_hr_SSIM = SSIM(self.pred_hr4,self.hr_data4)
        # self.loss_hr_col_SSIM = SSIM(self.pred_hr_col4,self.hr_data4)
        # self.loss_hr_ldiag_SSIM = SSIM(self.pred_ldiag_row4,self.hr_data4)
        # self.loss_hr_rdiag_SSIM = SSIM(self.pred_rdiag_col4,self.hr_data4)

        self.count_mean_hr_SSIM.update(self.loss_hr_SSIM)
        # self.count_mean_hr_col_SSIM.update(self.loss_hr_col_SSIM)
        # self.count_mean_hr_ldiag_SSIM.update(self.loss_hr_ldiag_SSIM)
        # self.count_mean_hr_rdiag_SSIM.update(self.loss_hr_rdiag_SSIM)

        self.loss_mean_hr_SSIM = self.count_mean_hr_SSIM.avg
        # self.loss_mean_hr_col_SSIM = self.count_mean_hr_col_SSIM.avg
        # self.loss_mean_hr_ldiag_SSIM = self.count_mean_hr_ldiag_SSIM.avg
        # self.loss_mean_hr_rdiag_SSIM = self.count_mean_hr_rdiag_SSIM.avg
        pass
    
    def iter_end(self):
        pass
    def epoch_end(self):
        if self.isTrain:
            # self.count_mean_LightField_row.reset()
            # self.count_mean_LightField_col.reset()
            # self.count_mean_LightField_ldiag.reset()
            # self.count_mean_LightField_rdiag.reset()
            self.count_mean_LightField.reset()
            pass
        else:
            self.count_mean_hr_PSNR.reset()
            # self.count_mean_hr_col_PSNR.reset()
            # self.count_mean_hr_ldiag_PSNR.reset()
            # self.count_mean_hr_rdiag_PSNR.reset()
            self.count_mean_hr_SSIM.reset()
            # self.count_mean_hr_col_SSIM.reset()
            # self.count_mean_hr_ldiag_SSIM.reset()
            # self.count_mean_hr_rdiag_SSIM.reset()
        pass
    def val_save(self):

        return False
        
        # if self.opt.isTrain:
        #     self.epoch+=1
            # self.fac = self.epoch/(self.nepoch+1)
